"""
创建 2段 四极杆质谱 电势阵列
结构: 地电极 + 4扇形电极 + 地电极 + 4扇形电极(旋转45°) + 地电极

参考 GEM 文件:
pa_define(200,200,1400,planar,non-mirrored,E,100,surface=fractional)
"""

import sys
import os
import math

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from SIMION.PA import PA


def create_quadrupole_ms(
    # 阵列尺寸
    nx=150,
    ny=150, 
    nz=1400,
    # 电极几何参数
    inner_radius=50,    # 内径
    outer_radius=60,    # 外径
    sector_angle=88,    # 扇形张角 (度)
    electrode_length=500,  # 扇形电极长度
    ring_thickness=10,  # 圆环电极厚度
    # 位置参数
    ground1_z=90,       # 第一个接地电极z位置
    quad1_z=101,        # 第一组四极起始z
    ground2_z=603,      # 中间接地电极z位置  
    quad2_z=614,        # 第二组四极起始z
    ground3_z=1116,     # 最后接地电极z位置
    # 角度参数
    angles1=None,       # 第一组扇形电极角度
    angles2=None,       # 第二组扇形电极角度 (旋转45°)
    # 输出文件
    output_file=None
):
    """
    创建双段四极杆结构
    
    结构说明:
    - 空间: nx x ny x nz (planar, 3D)
    - 中心: 自动计算 (nx//2, ny//2)
    
    电极布局 (沿z轴):
    1. 接地电极 (圆环) @ ground1_z
    2. 第一组4扇形电极 @ quad1_z, 长度 electrode_length
    3. 接地电极 (圆环) @ ground2_z
    4. 第二组4扇形电极 @ quad2_z (旋转45°)
    5. 接地电极 (圆环) @ ground3_z
    """
    
    # 中心坐标 (动态计算)
    cx, cy = nx // 2, ny // 2
    
    # 默认角度
    if angles1 is None:
        angles1 = [-44, 46, 136, 226]  # 对应 electrode 1,2,3,4
    if angles2 is None:
        angles2 = [1, 91, 181, 271]    # 对应 electrode 5,6,7,8
    
    # 圆环电极使用相同的内外径
    ring_inner = inner_radius
    ring_outer = outer_radius
    
    print("="*60)
    print("创建双段四极杆质谱电势阵列")
    print("="*60)
    print(f"\n阵列尺寸: {nx} x {ny} x {nz}")
    print(f"中心: ({cx}, {cy})")
    print(f"\n电极布局:")
    print(f"  接地电极1: z = {ground1_z} ~ {ground1_z + ring_thickness}")
    print(f"  四极组1:   z = {quad1_z} ~ {quad1_z + electrode_length}")
    print(f"  接地电极2: z = {ground2_z} ~ {ground2_z + ring_thickness}")
    print(f"  四极组2:   z = {quad2_z} ~ {quad2_z + electrode_length}")
    print(f"  接地电极3: z = {ground3_z} ~ {ground3_z + ring_thickness}")
    print(f"\n扇形参数:")
    print(f"  内径: {inner_radius}, 外径: {outer_radius}")
    print(f"  扇形角度: {sector_angle}°")
    print(f"  第一组角度: {angles1}")
    print(f"  第二组角度: {angles2} (旋转45°)")
    
    # 创建3D planar阵列
    pa = PA(
        nx=nx,
        ny=ny,
        nz=nz,
        symmetry='planar',
        mirror_x=0,
        mirror_y=0,
        mirror_z=0,
        max_voltage=100000.0,
        fast_adjustable=1
    )
    
    print(f"\n正在生成电极点...")
    
    def is_in_ring(x, y, r_inner, r_outer):
        """判断点是否在圆环内"""
        dx = x - cx
        dy = y - cy
        r = math.sqrt(dx*dx + dy*dy)
        return r_inner <= r <= r_outer
    
    def is_in_sector(x, y, center_angle, half_angle, r_inner, r_outer):
        """
        判断点是否在扇形内
        center_angle: 扇形中心角度 (度)
        half_angle: 扇形半角 (度)
        """
        dx = x - cx
        dy = y - cy
        r = math.sqrt(dx*dx + dy*dy)
        
        # 检查径向范围
        if not (r_inner <= r <= r_outer):
            return False
        
        # 计算角度 (度)
        angle = math.degrees(math.atan2(dy, dx))
        
        # 计算与中心角度的差值
        diff = angle - center_angle
        # 归一化到 -180 ~ 180
        while diff > 180:
            diff -= 360
        while diff < -180:
            diff += 360
        
        return abs(diff) <= half_angle
    
    half_sector = sector_angle / 2.0
    
    # 统计
    ground_count = 0
    sector_count = [0] * 8  # 8个扇形电极
    
    # 遍历所有点
    for z in range(nz):
        for y in range(ny):
            for x in range(nx):
                electrode_num = None
                voltage = None
                
                # 检查接地电极1
                if ground1_z <= z < ground1_z + ring_thickness:
                    if is_in_ring(x, y, ring_inner, ring_outer):
                        electrode_num = 0
                        voltage = 0.0
                        ground_count += 1
                
                # 检查第一组四极电极 (z = 101 ~ 601)
                elif quad1_z <= z < quad1_z + electrode_length:
                    for i, angle in enumerate(angles1):
                        if is_in_sector(x, y, angle, half_sector, inner_radius, outer_radius):
                            electrode_num = i + 1  # electrode 1,2,3,4
                            voltage = float(i + 1)
                            sector_count[i] += 1
                            break
                
                # 检查接地电极2
                elif ground2_z <= z < ground2_z + ring_thickness:
                    if is_in_ring(x, y, ring_inner, ring_outer):
                        electrode_num = 0
                        voltage = 0.0
                        ground_count += 1
                
                # 检查第二组四极电极 (z = 614 ~ 1114)
                elif quad2_z <= z < quad2_z + electrode_length:
                    for i, angle in enumerate(angles2):
                        if is_in_sector(x, y, angle, half_sector, inner_radius, outer_radius):
                            electrode_num = i + 5  # electrode 5,6,7,8
                            voltage = float(i + 5)
                            sector_count[i + 4] += 1
                            break
                
                # 检查接地电极3
                elif ground3_z <= z < ground3_z + ring_thickness:
                    if is_in_ring(x, y, ring_inner, ring_outer):
                        electrode_num = 0
                        voltage = 0.0
                        ground_count += 1
                
                # 设置电极点
                if electrode_num is not None:
                    pa.point(x, y, z, 1, voltage)
        
        # 进度显示
        if z % 100 == 0:
            print(f"  处理中... z = {z}/{nz}")
    
    print(f"\n电极点统计:")
    print(f"  接地电极: {ground_count} 点")
    for i in range(4):
        print(f"  扇形电极 {i+1}: {sector_count[i]} 点")
    for i in range(4, 8):
        print(f"  扇形电极 {i+1}: {sector_count[i]} 点")
    
    # 保存文件
    script_dir = os.path.dirname(os.path.abspath(__file__))
    if output_file is None:
        output_file = os.path.join(script_dir, f'quadrupole_ms_{nx}x{ny}x{nz}_{inner_radius * 2}Φ-{sector_angle}deg-{electrode_length}mm.pa#')
    elif not os.path.isabs(output_file):
        output_file = os.path.join(script_dir, output_file)
    
    print(f"\n保存文件: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("PA 文件头信息:")
    print("="*60)
    print(pa.header_string())
    
    print("\n电极编号说明:")
    print("  electrode(0): 接地电极 (3个圆环)")
    print("  electrode(1): 第一组扇形, 角度 -44°")
    print("  electrode(2): 第一组扇形, 角度 46°")
    print("  electrode(3): 第一组扇形, 角度 136°")
    print("  electrode(4): 第一组扇形, 角度 226°")
    print("  electrode(5): 第二组扇形, 角度 1°")
    print("  electrode(6): 第二组扇形, 角度 91°")
    print("  electrode(7): 第二组扇形, 角度 181°")
    print("  electrode(8): 第二组扇形, 角度 271°")
    
    return pa


if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='创建四极杆质谱电势阵列')
    parser.add_argument('--nx', type=int, default=150, help='X方向尺寸')
    parser.add_argument('--ny', type=int, default=150, help='Y方向尺寸')
    parser.add_argument('--nz', type=int, default=1400, help='Z方向尺寸')
    args = parser.parse_args()
    
    create_quadrupole_ms(nx=args.nx, ny=args.ny, nz=args.nz)

